package com.zt.chatgpt.openai.service.channel.impl;

import com.alibaba.fastjson.JSON;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.zt.chatglm.commons.Constants;
import com.zt.chatglm.domain.chat.ChatCompletionRequest;
import com.zt.chatglm.domain.chat.ChatCompletionResponse;
import com.zt.chatglm.session.OpenAiSession;
import com.zt.chatgpt.openai.model.aggregates.ChatProcessAggregate;
import com.zt.chatgpt.openai.model.entity.MessageEntity;
import com.zt.chatgpt.openai.service.channel.OpenAiGroupService;
import com.zt.chatgpt.types.enums.ChatGLMModel;
import com.zt.chatgpt.types.exception.ChatGPTException;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * @author: zt
 * @Description: ChatGLM 服务
 * @create: 2024-05-22 23:53
 */
@Slf4j
@Service
public class ChatGLMService implements OpenAiGroupService {

    @Resource
    protected OpenAiSession chatGLMOpenAiSession;

    @Override
    public void doMessageResponse(ChatProcessAggregate chatProcess, ResponseBodyEmitter emitter) throws Exception {
        // 1. 请求消息封装
        List<ChatCompletionRequest.Prompt> prompts = new ArrayList<>();

        List<MessageEntity> messages = chatProcess.getMessages();

        for (MessageEntity message : messages) {
            String role = message.getRole();
            if (Objects.equals(role, Constants.Role.system.getCode())) {
                prompts.add(ChatCompletionRequest.Prompt.builder()
                        .role(Constants.Role.system.getCode())
                        .content(message.getContent())
                        .build());
            } else {
                prompts.add(ChatCompletionRequest.Prompt.builder()
                        .role(Constants.Role.user.getCode())
                        .content(message.getContent())
                        .build());
            }
        }

        // 2. 构建请求
        ChatCompletionRequest request = new ChatCompletionRequest();
        request.setModel(Constants.Model.valueOf(ChatGLMModel.get(chatProcess.getModel()).name()));
        request.setPrompt(new ArrayList<ChatCompletionRequest.Prompt>(prompts));
        log.info("请求 for set {}", request.toString());
        ChatCompletionRequest request1 = ChatCompletionRequest.builder()
                .prompt(prompts)
                .model(Constants.Model.valueOf(ChatGLMModel.get(chatProcess.getModel()).name()))
                .build();
        log.info("请求 for builder {}", request1.toString());

        // 3. 发送请求
        chatGLMOpenAiSession.completions(request, new EventSourceListener() {
            @Override
            public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {

                ChatCompletionResponse response = JSON.parseObject(data, ChatCompletionResponse.class);

                // 发送信息
                if (Constants.EventType.add.getCode().equals(type)){
                    try {
                        emitter.send(response.getData());
                    } catch (Exception e) {
                        throw new ChatGPTException(e.getMessage());
                    }
                }

                // type 消息类型，add 增量，finish 结束，error 错误，interrupted 中断
                if (Constants.EventType.finish.getCode().equals(type)) {
                    ChatCompletionResponse.Meta meta = JSON.parseObject(response.getMeta(), ChatCompletionResponse.Meta.class);
                    log.info("[输出结束] Tokens {}", JSON.toJSONString(meta));
                }
            }
            @Override
            public void onClosed(@NotNull EventSource eventSource) {
                emitter.complete();
            }

            @Override
            public void onFailure(@NotNull EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                emitter.completeWithError(t);
                log.info("ChatGLM 异常：{}", t);
            }
        });

    }
}
